import datetime
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import logging
from typing import Dict, List, Optional, Any

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 尝试导入mplfinance，如果失败则设置为None
try:
    import mplfinance as mpf
except ImportError:
    mpf = None
    logger.warning("mplfinance not found. Candlestick charts will be disabled.")

from .data_handler import RealDataHandler, FakeDataHandler
from .indicators import calculate_sma, calculate_rsi, calculate_atr
from .portfolio import Portfolio
from .context import Context
from .exceptions import (
    BacktestError, DataError, ConfigurationError, 
    PortfolioError, StrategyError
)


class Broker:
    """交易执行器，负责处理订单和费用计算"""
    
    def __init__(self, portfolio: Portfolio, commission_info: Dict[str, float]):
        self.portfolio = portfolio
        self.commission_rate = commission_info.get('commission', 0.0003)
        self.stamp_duty_rate = commission_info.get('stamp_duty', 0.001)
        self.min_commission = commission_info.get('min_commission', 5.0)
        self.transaction_log = []
        
        logger.info(f"Broker initialized with commission rate: {self.commission_rate}")
    
    def execute_order(self, stock: str, shares: int, price: float, dt) -> bool:
        """执行订单"""
        try:
            if shares == 0:
                logger.warning(f"Order with 0 shares for {stock}, skipping")
                return False
            
            if price <= 0:
                logger.error(f"Invalid price {price} for {stock}")
                return False
            
            # 验证订单
            if not self.portfolio.validate_order(stock, shares, price):
                logger.error(f"Order validation failed for {stock}: {shares} shares @ {price}")
                return False
            
            # 计算交易成本
            trade_value = abs(shares) * price
            commission = max(trade_value * self.commission_rate, self.min_commission)
            stamp_duty = trade_value * self.stamp_duty_rate if shares < 0 else 0.0
            total_cost = commission + stamp_duty
            
            # 更新现金
            self.portfolio.cash -= (shares * price + total_cost)
            
            # 记录交易
            transaction = {
                'date': dt,
                'stock': stock,
                'type': 'BUY' if shares > 0 else 'SELL',
                'shares': abs(shares),
                'price': price,
                'value': trade_value,
                'cost': total_cost
            }
            self.transaction_log.append(transaction)
            
            # 更新持仓
            current_shares = self.portfolio.positions.get(stock, {}).get('shares', 0)
            new_shares = current_shares + shares
            
            if new_shares > 0:
                # 计算加权平均成本
                current_cost = self.portfolio.positions.get(stock, {}).get('cost_price', 0) * current_shares
                new_cost_price = (current_cost + shares * price) / new_shares
                
                self.portfolio.positions[stock] = {
                    'shares': new_shares,
                    'cost_price': new_cost_price,
                    'last_price': price
                }
            elif new_shares == 0:
                # 全部卖出，计算已实现损益
                if stock in self.portfolio.positions:
                    cost_price = self.portfolio.positions[stock]['cost_price']
                    realized_pnl = abs(shares) * (price - cost_price)
                    self.portfolio.add_realized_pnl(realized_pnl)
                    del self.portfolio.positions[stock]
            else:
                # 做空不允许
                logger.error(f"Short selling not allowed for {stock}")
                # 撤销现金更新
                self.portfolio.cash += (shares * price + total_cost)
                return False
            
            logger.info(f"  [Broker] Executed: {'BUY' if shares > 0 else 'SELL'} {abs(shares)} {stock} @ {price:.2f}")
            return True
            
        except Exception as e:
            logger.error(f"Error executing order for {stock}: {str(e)}")
            return False

class AnalysisResult:
    def __init__(self, equity_curve, dates, transactions, strategy_params, output_path, all_stock_data):
        self.equity_curve = pd.Series(equity_curve, index=pd.to_datetime(dates)).dropna()
        self.transactions = pd.DataFrame(transactions)
        self.strategy_params = strategy_params
        self.output_path = output_path
        self.all_stock_data = all_stock_data
        self.metrics = {}

    def _calculate_metrics(self):
        if self.equity_curve.empty: return
        self.metrics['initial_value'] = self.equity_curve.iloc[0]; self.metrics['final_value'] = self.equity_curve.iloc[-1]
        self.metrics['total_return_pct'] = (self.metrics['final_value'] / self.metrics['initial_value']) - 1
        daily_returns = self.equity_curve.pct_change().dropna()
        self.metrics['sharpe_ratio'] = (daily_returns.mean() / daily_returns.std()) * np.sqrt(252) if daily_returns.std() != 0 else 0.0
        self.metrics['max_drawdown_pct'] = (self.equity_curve - self.equity_curve.cummax()).div(self.equity_curve.cummax()).min()
        if not self.transactions.empty:
            self.metrics.update({'total_trades': len(self.transactions), 'total_commission': self.transactions['cost'].sum(),
                                 'buy_trades': len(self.transactions[self.transactions['type'] == 'BUY']),
                                 'sell_trades': len(self.transactions[self.transactions['type'] == 'SELL'])})
        else: self.metrics.update({'total_trades': 0, 'total_commission': 0, 'buy_trades': 0, 'sell_trades': 0})

    def plot_equity_curve(self, show=True, save=True):
        """绘制权益曲线图"""
        if self.equity_curve.empty: 
            print("[Analysis] No trades were made. Cannot generate equity curve plot.")
            return
        
        try:
            plt.style.use('seaborn-v0_8-darkgrid')
        except OSError:
            # 如果上述样式不可用，使用默认样式
            pass
            
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(16, 10), 
                                       gridspec_kw={'height_ratios': [3, 1]}, 
                                       sharex=True)
        
        ax1.plot(self.equity_curve.index, self.equity_curve, 
                label='Strategy Equity', color='royalblue', linewidth=2)
        ax1.set_title(f'Strategy Performance Analysis (Run ID: {os.path.basename(self.output_path)})', 
                     fontsize=16)
        ax1.set_ylabel('Total Value')
        ax1.legend()
        ax1.grid(True)
        
        drawdown = (self.equity_curve - self.equity_curve.cummax()) / self.equity_curve.cummax()
        ax2.fill_between(drawdown.index, drawdown, 0, color='red', alpha=0.3)
        ax2.set_ylabel('Drawdown')
        ax2.set_xlabel('Date')
        ax2.grid(True)
        
        plt.tight_layout()
        
        if save:
            file_path = os.path.join(self.output_path, 'performance_chart.png')
            plt.savefig(file_path)
            print(f"[Analysis] Performance chart saved to: {file_path}")
        
        if show:
            plt.show()
        else:
            plt.close(fig)

    def plot_trades_on_candlestick(self, save=True):
        """在K线图上显示交易信号"""
        if not self.all_stock_data:
            print("[Analysis] No stock data available to plot.")
            return
        
        if mpf is None:
            print("[Analysis] mplfinance not available. Cannot generate candlestick charts.")
            return
            
        charts_path = os.path.join(self.output_path, 'trade_charts')
        os.makedirs(charts_path, exist_ok=True)
        print(f"[Analysis] Generating candlestick dashboards in: '{charts_path}'")

        for stock, stock_data in self.all_stock_data.items():
            try:
                if stock_data is None or len(stock_data) < self.strategy_params.get('sma_long_window', 20):
                    print(f"  [Warning] Insufficient data for {stock}, skipping.")
                    continue
                
                p = self.strategy_params
                sma_long_window = p.get('sma_long_window', 20)
                rsi_window = p.get('rsi_window', 14)
                atr_window = p.get('atr_window', 14)
                
                sma_long_series = calculate_sma(stock_data['close'], sma_long_window, return_series=True)
                rsi_series = calculate_rsi(stock_data['close'], rsi_window, return_series=True)
                atr_series = calculate_atr(stock_data['high'], stock_data['low'], 
                                         stock_data['close'], atr_window, return_series=True)

                add_plots = [
                    mpf.make_addplot(sma_long_series, panel=0, color='blue', width=0.7),
                    mpf.make_addplot(rsi_series, panel=2, color='purple', ylabel='RSI'),
                    mpf.make_addplot(atr_series, panel=3, color='orange', ylabel='ATR')
                ]
                plot_title = f'Analysis Dashboard for {stock}'
                
                if not self.transactions.empty and stock in self.transactions['stock'].values:
                    plot_title = f'Trades on {stock}'
                    stock_trans = self.transactions[self.transactions['stock'] == stock]
                    buy_signals = stock_trans[stock_trans['type'] == 'BUY']
                    sell_signals = stock_trans[stock_trans['type'] == 'SELL']
                    
                    if not buy_signals.empty:
                        buy_markers = pd.Series(np.nan, index=stock_data.index)
                        buy_markers.loc[buy_signals['date']] = buy_signals['price'] * 0.98
                        add_plots.append(mpf.make_addplot(buy_markers, type='scatter', 
                                                         panel=0, marker='^', 
                                                         color='lime', markersize=120))
                    
                    if not sell_signals.empty:
                        sell_markers = pd.Series(np.nan, index=stock_data.index)
                        sell_markers.loc[sell_signals['date']] = sell_signals['price'] * 1.02
                        add_plots.append(mpf.make_addplot(sell_markers, type='scatter', 
                                                         panel=0, marker='v', 
                                                         color='red', markersize=120))

                chart_file = os.path.join(charts_path, f"dashboard_{stock}.png")
                
                hlines_dict = {}
                if 'rsi_pullback_level' in p and 'rsi_profit_target' in p:
                    hlines_dict = dict(
                        hlines=[p['rsi_pullback_level'], p['rsi_profit_target']], 
                        colors=['g', 'r'], 
                        linestyle='--'
                    )
                
                mpf.plot(
                    stock_data, type='candle', style='charles', title=plot_title,
                    ylabel='Price (Panel 0)', volume=True, addplot=add_plots,
                    panel_ratios=(6, 1, 2, 2),
                    hlines=hlines_dict,
                    savefig=chart_file, figsize=(19, 10), 
                    warn_too_much_data=len(stock_data) + 1
                )
                print(f"  - Dashboard for {stock} saved.")
                
            except Exception as e:
                print(f"  [Error] Failed to generate chart for {stock}: {str(e)}")
                continue

    def generate_report(self):
        self._calculate_metrics()
        report_path = os.path.join(self.output_path, 'report.txt')
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write("="*57 + "\n" + "           DragonQuant Backtest Analysis Report\n" + "="*57 + "\n\n")
            f.write(f"Run ID: {os.path.basename(self.output_path)}\nReport Generated: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            f.write("--- Strategy Parameters ---\n")
            for key, value in self.strategy_params.items(): f.write(f"{key:<25}: {value}\n")
            f.write("\n")
            if self.equity_curve.empty: f.write("--- Result ---\nNo trades were executed.\n"); return
            f.write("--- Overall Performance ---\n")
            f.write(f"{'Initial Value':<25}: {self.metrics.get('initial_value', 0):,.2f}\n")
            f.write(f"{'Final Value':<25}: {self.metrics.get('final_value', 0):,.2f}\n")
            f.write(f"{'Total Return':<25}: {self.metrics.get('total_return_pct', 0):.2%}\n\n")
            f.write("--- Risk & Return Metrics ---\n")
            f.write(f"{'Sharpe Ratio (Annualized)':<25}: {self.metrics.get('sharpe_ratio', 0):.3f}\n")
            f.write(f"{'Max Drawdown':<25}: {self.metrics.get('max_drawdown_pct', 0):.2%}\n\n")
            f.write("--- Trading Activity ---\n")
            f.write(f"{'Total Trades':<25}: {self.metrics.get('total_trades', 0)}\n")
            f.write(f"{'Buy Transactions':<25}: {self.metrics.get('buy_trades', 0)}\n")
            f.write(f"{'Sell Transactions':<25}: {self.metrics.get('sell_trades', 0)}\n")
            f.write(f"{'Total Commission & Fees':<25}: {self.metrics.get('total_commission', 0):,.2f}\n\n")
        print(f"[Analysis] Detailed report saved to: {report_path}")
    
    def plot(self, show=True, save=True):
        """主要的绘图方法，同时显示权益曲线和交易图表"""
        self.plot_equity_curve(show=show, save=save)
        self.plot_trades_on_candlestick(save=save)

class Cerebro:
    """回测引擎核心类"""
    
    def __init__(self):
        self._strategy_cls = None
        self.universe = []
        self._config = {
            'start_date': None,
            'end_date': None,
            'initial_cash': 1e6,
            'data_source': 'fake'
        }
        self._commission_info = {
            'commission': 0.0003,
            'stamp_duty': 0.001,
            'min_commission': 5.0
        }
        self._orders_to_execute = []
        
        logger.info("Cerebro backtest engine initialized")

    def add_strategy(self, strategy_cls):
        """添加策略类"""
        try:
            if not hasattr(strategy_cls, 'next'):
                raise StrategyError(f"Strategy class {strategy_cls.__name__} must implement 'next' method")
            
            self._strategy_cls = strategy_cls
            logger.info(f"Strategy '{strategy_cls.__name__}' loaded successfully")
            
        except Exception as e:
            logger.error(f"Failed to add strategy: {str(e)}")
            raise StrategyError(f"Failed to add strategy: {str(e)}")

    def set_universe(self, universe: List[str]):
        """设置股票池"""
        try:
            if not universe or len(universe) == 0:
                raise ConfigurationError("Universe cannot be empty")
            
            if not isinstance(universe, list):
                raise ConfigurationError("Universe must be a list of stock symbols")
            
            self.universe = universe
            logger.info(f"Universe set with {len(universe)} stocks: {universe[:5]}{'...' if len(universe) > 5 else ''}")
            
        except Exception as e:
            logger.error(f"Failed to set universe: {str(e)}")
            raise ConfigurationError(f"Failed to set universe: {str(e)}")

    def set_commission(self, commission: float = 0.0003, stamp_duty: float = 0.001, min_commission: float = 5.0):
        """设置交易费用"""
        try:
            if commission < 0 or stamp_duty < 0 or min_commission < 0:
                raise ConfigurationError("Commission rates cannot be negative")
            
            self._commission_info.update({
                'commission': commission,
                'stamp_duty': stamp_duty,
                'min_commission': min_commission
            })
            
            logger.info(f"Commission set: Rate={commission}, Stamp Duty={stamp_duty}, Min={min_commission}")
            
        except Exception as e:
            logger.error(f"Failed to set commission: {str(e)}")
            raise ConfigurationError(f"Failed to set commission: {str(e)}")

    def configure(self, start_date: str, end_date: str, initial_cash: float, 
                  data_source: str = 'fake', data_path: Optional[str] = None, **kwargs):
        """配置回测参数"""
        try:
            # 验证日期
            start_dt = pd.to_datetime(start_date)
            end_dt = pd.to_datetime(end_date)
            
            if start_dt >= end_dt:
                raise ConfigurationError("Start date must be before end date")
            
            if initial_cash <= 0:
                raise ConfigurationError("Initial cash must be positive")
            
            if data_source not in ['fake', 'real']:
                raise ConfigurationError("Data source must be 'fake' or 'real'")
            
            if data_source == 'real' and not data_path:
                raise ConfigurationError("Data path is required for real data source")
            
            self._config.update({
                'start_date': start_date,
                'end_date': end_date,
                'initial_cash': initial_cash,
                'data_source': data_source,
                'data_path': data_path
            })
            
            logger.info(f"Configuration complete: {start_date} to {end_date}, Cash: {initial_cash:,.0f}, Source: {data_source}")
            
        except Exception as e:
            logger.error(f"Configuration failed: {str(e)}")
            raise ConfigurationError(f"Configuration failed: {str(e)}")
    
    def _validate_configuration(self):
        """验证配置"""
        if not self._strategy_cls:
            raise ConfigurationError("No strategy has been added")
        
        if not self.universe:
            raise ConfigurationError("Universe has not been set")
        
        required_configs = ['start_date', 'end_date', 'initial_cash']
        missing_configs = [key for key in required_configs if self._config[key] is None]
        
        if missing_configs:
            raise ConfigurationError(f"Missing required configurations: {missing_configs}")
    
    def _create_order(self, stock: str, shares: int, style: str = 'amount'):
        """创建订单"""
        try:
            if style == 'target':
                current_shares = self._portfolio.positions.get(stock, {}).get('shares', 0)
                shares_to_order = shares - current_shares
            else:
                shares_to_order = shares
            
            if shares_to_order != 0:
                self._orders_to_execute.append({
                    'stock': stock,
                    'shares': shares_to_order
                })
                
        except Exception as e:
            logger.error(f"Failed to create order for {stock}: {str(e)}")
    
    def run(self) -> Optional['AnalysisResult']:
        """运行回测"""
        try:
            # 验证配置
            self._validate_configuration()
            
            # 创建输出目录
            run_id = f"run_{time.strftime('%Y%m%d_%H%M%S')}"
            output_path = os.path.join('backtest_results', run_id)
            os.makedirs(output_path, exist_ok=True)
            logger.info(f"Results will be saved in: '{output_path}'")

            logger.info("--- Initializing Backtest Engine ---")
            
            # 初始化策略
            try:
                strategy = self._strategy_cls()
                strategy.universe = self.universe
            except Exception as e:
                raise StrategyError(f"Failed to initialize strategy: {str(e)}")
            
            # 初始化投资组合和经纪人
            self._portfolio = Portfolio(self._config['initial_cash'])
            broker = Broker(self._portfolio, self._commission_info)
            
            # 设置API中的cerebro实例
            try:
                import dragonquant.api
                dragonquant.api.set_cerebro(self)
            except ImportError:
                logger.warning("Could not import dragonquant.api")

            # 初始化数据处理器
            try:
                if self._config['data_source'] == 'real':
                    data_handler = RealDataHandler(self.universe, self._config['data_path'])
                    # 获取所有数据的日期范围
                    all_dates = pd.DatetimeIndex([])
                    for df in data_handler.all_data.values():
                        if not df.empty:
                            all_dates = all_dates.union(df.index)
                    
                    if all_dates.empty:
                        raise DataError("No valid data found in data files")
                    
                    start_date = pd.to_datetime(self._config['start_date'])
                    end_date = pd.to_datetime(self._config['end_date'])
                    dates_to_run = all_dates[(all_dates >= start_date) & (all_dates <= end_date)]
                    
                    if dates_to_run.empty:
                        raise DataError(f"No data available for the specified date range: {start_date} to {end_date}")
                        
                else:
                    data_handler = FakeDataHandler(self.universe, self._config['start_date'], self._config['end_date'])
                    dates_to_run = data_handler.dates
                    
            except Exception as e:
                raise DataError(f"Failed to initialize data handler: {str(e)}")
            
            # 初始化策略
            try:
                strategy.init()
                logger.info("Strategy initialized successfully")
            except Exception as e:
                raise StrategyError(f"Strategy initialization failed: {str(e)}")
            
            logger.info(f"--- Starting Event Loop ({len(dates_to_run)} trading days) ---")
            equity_curve_dates = []
            processed_days = 0
            
            for dt in dates_to_run:
                try:
                    logger.debug(f"Processing Date: {dt.strftime('%Y-%m-%d')}")
                    
                    # 创建上下文
                    context = Context(dt, self.universe, data_handler, self._portfolio)
                    
                    # 执行策略
                    try:
                        strategy.next(context)
                    except Exception as e:
                        logger.error(f"Strategy execution failed on {dt}: {str(e)}")
                        continue
                    
                    # 获取当前价格
                    current_prices = data_handler.get_current_prices(dt)
                    if not current_prices:
                        logger.warning(f"No current prices available for {dt}")
                        continue
                    
                    # 执行订单
                    executed_orders = 0
                    for order in self._orders_to_execute:
                        stock = order['stock']
                        shares = order['shares']
                        
                        if stock in current_prices:
                            price = current_prices[stock]
                            if broker.execute_order(stock, shares, price, dt):
                                executed_orders += 1
                        else:
                            logger.warning(f"No price available for {stock} on {dt}")
                    
                    self._orders_to_execute.clear()
                    
                    # 更新投资组合价值
                    total_value = self._portfolio.update_total_value(current_prices)
                    self._portfolio.equity_curve.append(total_value)
                    equity_curve_dates.append(dt)
                    
                    processed_days += 1
                    
                    if processed_days % 50 == 0:
                        logger.info(f"Processed {processed_days}/{len(dates_to_run)} days, Current Equity: {total_value:,.2f}")
                    
                except Exception as e:
                    logger.error(f"Error processing date {dt}: {str(e)}")
                    continue
            
            logger.info(f"--- Backtest Finished (Processed {processed_days} days) ---")
            
            # 创建结果对象
            try:
                # 获取策略参数
                strategy_params = getattr(strategy, '__dict__', {})
                # 过滤掉非参数属性
                strategy_params = {k: v for k, v in strategy_params.items() 
                                 if not k.startswith('_') and not callable(v)}
                
                results = AnalysisResult(
                    equity_curve=self._portfolio.equity_curve,
                    dates=equity_curve_dates,
                    transactions=broker.transaction_log,
                    strategy_params=strategy_params,
                    output_path=output_path,
                    all_stock_data=data_handler.all_data
                )
                
                logger.info(f"Backtest completed successfully. Results available in AnalysisResult object.")
                return results
                
            except Exception as e:
                logger.error(f"Failed to create analysis results: {str(e)}")
                raise BacktestError(f"Failed to create analysis results: {str(e)}")
            
        except Exception as e:
            logger.error(f"Backtest failed: {str(e)}")
            raise BacktestError(f"Backtest failed: {str(e)}")